﻿using System;
using System.Collections.Generic;
using System.Linq;
using System.Text;

namespace LambdaExtensions.TreeExtensions
{
    /// <summary>
    /// Extension class for traversing through trees
    /// </summary>
    public static class TraversingExtensions
    {
        public static IEnumerable<T> SelectTreeManyDeep<T>(this IEnumerable<T> rootNodes, Func<T, IEnumerable<T>> childrenSelector)
        {
            Stack<T> Nodes = new Stack<T>();
            HashSet<T> Visited = new HashSet<T>();

            foreach (var Root in rootNodes)
                Nodes.Push(Root);

            while (Nodes.Any())
            {
                var Current = Nodes.Pop();

                if (!Visited.Add(Current))
                    continue;

                yield return Current;

                var children = childrenSelector(Current);
                if (children != null)
                    foreach (var c in children)
                    Nodes.Push(c);
            }
        }


        /// <summary>
        /// traverses a tree beginning from a set of root nodes, e.g. TreeView.Nodes 
        /// </summary>
        /// <typeparam name="T">Type of the child nodes</typeparam>
        /// <param name="rootNodes">List of root nodes, e.g. TreeView.Nodes </param>
        /// <param name="childrenSelector">selector returning of a node all sub nodes</param>
        /// <param name="TraverseAction">Action performed on every root and child node</param>
        public static IEnumerable<T> SelectTreeManyWide<T>(this IEnumerable<T> rootNodes, Func<T, IEnumerable<T>> childrenSelector)
        {
            Queue<T> Nodes = new Queue<T>();
            HashSet<T> Visited = new HashSet<T>();

            foreach (var Root in rootNodes)
                Nodes.Enqueue(Root);

            while (Nodes.Any())
            {
                var Current = Nodes.Dequeue();

                if (!Visited.Add(Current))
                    continue;

                yield return Current;

                var children = childrenSelector(Current);
                if (children != null)
                    foreach (var c in children)
                    Nodes.Enqueue(c);
            }
        }

        
        public static IEnumerable<T> SelectTreeDeep<T>(this T root, Func<T, IEnumerable<T>> childrenSelector)
        {
            return SelectTreeManyDeep(new[] { root }, childrenSelector);
        }

        public static IEnumerable<T> SelectTreeWide<T>(this T root, Func<T, IEnumerable<T>> childrenSelector)
        {
            return SelectTreeManyWide(new[] { root }, childrenSelector);
        }


        /// <summary>
        /// traverses a tree back to the root
        /// </summary>
        /// <typeparam name="T">Type of the child nodes</typeparam>
        /// <param name="child"></param>
        /// <param name="parentSelector">shall return the parent element, or null or this if the node is the root node</param>
        /// <param name="TraverseAction">Action performed on every root and child node</param>
        /// <returns>returns the root node</returns>
        public static IEnumerable<T> SelectTreeBack<T>(this T child, Func<T, T> parentSelector) where T : class
        {
            var CurrentNode = child;
            T last = null;

            HashSet<T> Visited = new HashSet<T>();

            while (child != null && child != last)
            {
                if (!Visited.Add(CurrentNode))
                    throw new ArgumentException("Circle detected");


                yield return child;

                last = CurrentNode;
                CurrentNode = parentSelector(child);
            }
        }

    

    }
}
